package com.chatplus.application.listener;

import cn.bugstack.openai.session.Configuration;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.chatplus.application.aiprocessor.platform.image.mj.MjImageProcessor;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.OkHttpClientUtil;
import com.chatplus.application.common.util.PlusJsonUtils;
import com.chatplus.application.domain.dto.ApiKeyDto;
import com.chatplus.application.domain.entity.draw.MjJobEntity;
import com.chatplus.application.domain.notification.MJJobNotification;
import com.chatplus.application.enumeration.AiPlatformEnum;
import com.chatplus.application.enumeration.ImageTaskTypeEnum;
import com.chatplus.application.service.account.UserProductLogService;
import com.chatplus.application.service.draw.MjJobService;
import com.chatplus.application.util.ConfigUtil;
import com.chatplus.application.web.notification.NotificationListener;
import okhttp3.*;
import org.springframework.stereotype.Component;

import java.io.IOException;

/**
 * MJ 任务监听
 */
@Component
public class MJJobNotificationListener {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(MJJobNotificationListener.class);
    private final MjJobService mjJobService;
    private final UserProductLogService userProductLogService;
    private final MjImageProcessor mjImageProcessor;

    public MJJobNotificationListener(MjJobService mjJobService,
                                     UserProductLogService userProductLogService,
                                     MjImageProcessor mjImageProcessor) {
        this.mjJobService = mjJobService;
        this.userProductLogService = userProductLogService;
        this.mjImageProcessor = mjImageProcessor;
    }

    @NotificationListener(notification = MJJobNotification.class)
    public void handleMJJobNotification(MJJobNotification notification) {
        if (notification == null) {
            return;
        }
        Long mjJobId = notification.getMjJobId();
        LOGGER.message("接收到MJ请求通知").context("mjJobId", mjJobId).info();
        MjJobEntity mjJobEntity = mjJobService.getById(mjJobId);
        if (mjJobEntity == null) {
            LOGGER.message("MJ任务不存在").context("mjJobId", mjJobId).error();
            return;
        }
        ApiKeyDto apiKeyDto = notification.getApiKeyDto();
        OkHttpClient okHttpClient = OkHttpClientUtil.getOkHttpClient(apiKeyDto.getProxyUrl(), true);
        Request request = new Request.Builder()
                .addHeader("Content-Type", Configuration.APPLICATION_JSON)
                .addHeader("Accept", Configuration.APPLICATION_JSON)
                .addHeader("Authorization", "Bearer " + apiKeyDto.getApiKey())
                .url(apiKeyDto.getUrl() + getUrl(mjJobEntity))
                .post(RequestBody.create(PlusJsonUtils.toJsonString(notification.getMjJobDto())
                        , MediaType.parse(Configuration.APPLICATION_JSON)))
                .build();
        Call call = okHttpClient.newCall(request);
        try (Response response = call.execute()) {
            if (response.isSuccessful() && response.body() != null) {
                String body = response.body().string();
                LOGGER.message("成功发送MJ图片生成请求").context("response", response).context("body", body).info();
                // 1(提交成功), 21(已存在), 22(排队中), other(错误)
                JSONObject resultJson = JSONUtil.parseObj(body);
                String code = resultJson.getStr("code");
                String description = resultJson.getStr("description");
                mjJobEntity.setErrMsg(description);
                if ("1".equals(code) || "22".equals(code)) {
                    String result = resultJson.getStr("result");
                    mjJobEntity.setTaskId(result);
                    mjJobEntity.setProgress(0);
                    JSONObject properties = resultJson.getJSONObject("properties");
                    mjJobEntity.setChannelId(properties.getStr("discordChannelId"));
                    userProductLogService.reducePower(mjJobEntity.getUserId(), ConfigUtil.getImagePower(mjJobEntity.getUserId(), AiPlatformEnum.MJ, mjJobEntity.getType()), AiPlatformEnum.MJ, null);
                } else {
                    mjJobEntity.setProgress(-1);
                    LOGGER.message("MJ图片生成失败").context("response", response).context("body", body).error();
                }
            } else {
                mjJobEntity.setProgress(-1);
                LOGGER.message("MJ图片请求失败").context("response", response).error();
            }
        } catch (IOException e) {
            mjJobEntity.setErrMsg("MJ图片请求失败");
            mjJobEntity.setProgress(-1);
            LOGGER.message("MJ图片请求失败").context("notification", notification).exception(e).error();
        } finally {
            mjJobService.updateById(mjJobEntity);
            mjImageProcessor.notifyUpdateTask(mjJobEntity.getId());
        }
    }

    private String getUrl(MjJobEntity mjJobEntity) {
        ImageTaskTypeEnum taskTypeEnum = mjJobEntity.getType();
        return switch (taskTypeEnum) {
            case TEXT_2_IMG -> "/mj/submit/imagine";
            case BLEND -> "/mj/submit/blend";
            case SWAP_FACE -> "/mj/insight-face/swap";
            case UPSCALE, VARIATION, ACTION -> "/mj/submit/action";
        };
    }
}
